import json
import os

import numpy

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

import operator

from torch.utils.data import Dataset
import pandas as pd
import torch
import torch.nn as nn
import numpy as np
from vedo import *
from scipy.spatial import distance_matrix

from utils import color_2_label


class Mesh_Dataset(Dataset):
    def __init__(self, data_list_path, num_classes=15, patch_size=6000):
        self.data_list = pd.read_csv(data_list_path, header=None)
        self.num_classes = num_classes
        self.patch_size = patch_size

    def __len__(self):
        return self.data_list.shape[0]

    def __getitem__(self, idx):
        with open('config.json', 'r') as f:
            config = json.load(f)

        mod = config["mod"]
        dataset_path = config["data"]["dataset_path"]
        label_path = dataset_path + config["data"]["label_data_path"] + mod

        with open(os.path.join(label_path, 'label.json'), 'r') as f:
            labels_cls_map = json.load(f)
        with open(os.path.join(label_path, 'edge_index.json'), 'r') as f:
            edge_index_map = json.load(f)

        if torch.is_tensor(idx):
            idx = idx.tolist()

        i_mesh = self.data_list.iloc[idx][0]
        label_i_mesh = os.path.join(label_path, i_mesh.split('_')[-1])

        mesh = Mesh(i_mesh)
        label_mesh = Mesh(label_i_mesh)

        colors = mesh.cellcolors[:, 0:3] / 255.
        label_colors = label_mesh.cellcolors
        labels_seg = color_2_label(label_colors, mod=mod)
        labels_cls = labels_cls_map[i_mesh.split('_')[-1].split('.')[0]]
        labels_cls = np.asarray(labels_cls).astype(dtype='int64')

        edge_index = edge_index_map[i_mesh.split('_')[-1].split('.')[0]]
        edge_index = np.asarray(edge_index).astype(dtype='int32')

        A = torch.from_numpy(np.eye(15, dtype='float'))
        for (x, y) in edge_index:
            A[x][y] = 1
            A[y][x] = 1
        d = A.sum(1)
        D = torch.diag(torch.pow(d, -0.5))
        normalized_adj = D.mm(A).mm(D)

        points = mesh.points()
        mean_cell_centers = mesh.center_of_mass()
        points[:, 0:3] -= mean_cell_centers[0:3]

        ids = np.array(mesh.faces())
        cells = points[ids].reshape(mesh.ncells, 9).astype(dtype='float32')

        mesh.compute_normals()
        normals = mesh.celldata['Normals']

        barycenters = mesh.cell_centers()
        barycenters -= mean_cell_centers[0:3]

        maxs = points.max(axis=0)
        mins = points.min(axis=0)
        means = points.mean(axis=0)
        stds = points.std(axis=0)
        nmeans = normals.mean(axis=0)
        nstds = normals.std(axis=0)

        for i in range(3):
            cells[:, i] = (cells[:, i] - means[i]) / stds[i]
            cells[:, i + 3] = (cells[:, i + 3] - means[i]) / stds[i]
            cells[:, i + 6] = (cells[:, i + 6] - means[i]) / stds[i]
            barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
            normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]

        # X = np.column_stack((cells, barycenters, normals, colors))
        X = np.column_stack((cells, barycenters, normals))
        Y = labels_seg

        X_train = np.zeros([self.patch_size, X.shape[1]], dtype='float32')
        Y_train = np.zeros([self.patch_size, Y.shape[1]], dtype='int32')
        S1 = np.zeros([self.patch_size, self.patch_size], dtype='float32')
        S2 = np.zeros([self.patch_size, self.patch_size], dtype='float32')

        # positive_idx = np.argwhere(labels_seg > 0)[:, 0]  # tooth idx
        # negative_idx = np.argwhere(labels_seg == 0)[:, 0]  # gingiva idx

        # num_positive = len(positive_idx)  # number of selected tooth cells
        #
        # if num_positive > self.patch_size:  # all positive_idx in this patch
        #     positive_selected_idx = np.random.choice(positive_idx, size=self.patch_size, replace=False)
        #     selected_idx = positive_selected_idx
        # else:  # patch contains all positive_idx and some negative_idx
        #     num_negative = self.patch_size - num_positive  # number of selected gingiva cells
        #     positive_selected_idx = np.random.choice(positive_idx, size=num_positive, replace=False)
        #     negative_selected_idx = np.random.choice(negative_idx, size=num_negative, replace=False)
        #     selected_idx = np.concatenate((positive_selected_idx, negative_selected_idx))
        #
        # selected_idx = np.sort(selected_idx, axis=None)

        valid_idx = np.argwhere(labels_seg >= 0)[:, 0]
        print(valid_idx.shape)
        print(valid_idx)
        selected_idx = np.sort(np.random.choice(valid_idx, size=self.patch_size, replace=False))

        X_train[:] = X[selected_idx, :]
        Y_train[:] = Y[selected_idx, :]

        if torch.cuda.is_available():
            TX = torch.as_tensor(X_train[:, 9:12], device='cuda')
            TD = torch.cdist(TX, TX)
            D = TD.cpu().numpy()
        else:
            D = distance_matrix(X_train[:, 9:12], X_train[:, 9:12])

        S1[D<0.1] = 1.0
        S1 = S1 / np.dot(np.sum(S1, axis=1, keepdims=True), np.ones((1, self.patch_size)))

        S2[D<0.2] = 1.0
        S2 = S2 / np.dot(np.sum(S2, axis=1, keepdims=True), np.ones((1, self.patch_size)))

        X_train = X_train.transpose(1, 0)
        Y_train = Y_train.transpose(1, 0)

        sample = {
            'cells': torch.from_numpy(X_train),
            'labels_seg': torch.from_numpy(Y_train),
            'labels_cls': torch.from_numpy(labels_cls),
            'A_S': torch.from_numpy(S1),
            'A_L': torch.from_numpy(S2),
            'normalized_adj': normalized_adj
        }

        return sample


if __name__ == '__main__':
    dataset = Mesh_Dataset('./list_lower/train_list_1.csv')
    sample = dataset.__getitem__(0)
    cells = sample["cells"]
    labels_seg = sample["labels_seg"]
    labels_cls = sample["labels_cls"]
    A_S = sample["A_S"]
    A_L = sample["A_L"]
    normalized_adj = sample["normalized_adj"]

    # print(cells)
    # print(labels_seg)
    # print(labels_cls)
    # print(nn.functional.one_hot(labels_cls, num_classes=5))
    # print(A_S)
    # print(A_L)
    # print(normalized_adj)